import queue
import threading

from   vllm.logger import init_logger

logger = init_logger(__name__)


class TaskQueue:
    def __init__(self):
        self._queue         = queue.Queue()
        self._worker_thread = threading.Thread(target=self._process_queue)
        self._running       = False
        self._start_lock    = threading.Lock()

    def enqueue(self, func, *args, **kwargs):
        """将函数及其参数放入队列。"""
        self._queue.put((func, args, kwargs))

    def _process_queue(self):
        """从队列中取出函数并执行。"""
        while self._running:
            func, args, kwargs = self._queue.get()
            if func is None:
                self._queue.task_done()
                break
            try:
                func(*args, **kwargs)
            except Exception as e:
                logger.warning(f"执行函数 {func.__name__} 时发生错误: {e}")
            finally:
                #! 队列中的计数减1，计数为0时表示任务被完成，join将不再阻塞
                self._queue.task_done()  # 标记任务完成，允许 join() 阻塞

    def start_worker(self):
        """启动工作线程（如果尚未启动）。"""
        with self._start_lock:
            if not self._running:
                self._running = True
                self._worker_thread.start()

    def join(self):
        """阻塞直到队列中的所有任务都被处理完毕。"""
        self._queue.join()

    def stop(self):
        """停止工作线程。这会在当前正在执行的任务完成后停止。"""
        with self._start_lock:
            if self._running:
                self._running = False
                self._queue.put((None, None, None))
                self._worker_thread.join()